Skip to content

build(deps-dev): bump jax from 0.5.0 to 0.6.2#5102

Draft
dependabot[bot] wants to merge 1 commit into
masterfrom
dependabot/pip/devel/jax-0.6.2
Draft

build(deps-dev): bump jax from 0.5.0 to 0.6.2#5102
dependabot[bot] wants to merge 1 commit into
masterfrom
dependabot/pip/devel/jax-0.6.2

Conversation

@dependabot
Copy link
Copy Markdown
Contributor

@dependabot dependabot Bot commented on behalf of github Dec 15, 2025 •
edited
Loading

Bumps jax from 0.5.0 to 0.6.2.

Release notes

Sourced from jax's releases.

JAX v0.6.2

  • New features:

    • Added jax.tree.broadcast which implements a pytree prefix broadcasting helper.

  • Changes

    • The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12.

JAX v0.6.1

  • New features:

    • Added jax.lax.axis_size which returns the size of the mapped axis given its name.

  • Changes

    • Additional checking for the versions of CUDA package dependencies was reenabled, having been accidentally disabled in a previous release.
    • JAX nightly packages are now published to artifact registry. To install these packages, see the JAX installation guide.
    • jax.sharding.PartitionSpec no longer inherits from a tuple.
    • jax.ShapeDtypeStruct is immutable now. Please use .update method to update your ShapeDtypeStruct instead of doing in-place updates.

  • Deprecations

    • jax.custom_derivatives.custom_jvp_call_jaxpr_p is deprecated, and will be removed in JAX v0.7.0.

JAX v0.6.0

  • Breaking changes

    • jax.numpy.array no longer accepts None. This behavior was deprecated since November 2023 and is now removed.
    • Removed the config.jax_data_dependent_tracing_fallback config option, which was added temporarily in v0.4.36 to allow users to opt out of the new "stackless" tracing machinery.
    • Removed the config.jax_eager_pmap config option.
    • Disallow the calling of lower and trace AOT APIs on the result of jax.jit if there have been subsequent wrappers applied. Previously this worked, but silently ignored the wrappers. The workaround is to apply jax.jit last among the wrappers, and similarly for jax.pmap. See [#27873](https://github.com/jax-ml/jax/issues/27873).
    • The cuda12_pip extra for jax has been removed; use pip install jax[cuda12] instead.

  • Changes

    • The minimum CuDNN version is v9.8.
    • JAX is now built using CUDA 12.8. All versions of CUDA 12.1 or newer remain supported.
    • JAX package extras are now updated to use dash instead of underscore to align with PEP 685. For instance, if you were previously using pip install jax[cuda12_local]

... (truncated)

Changelog

Sourced from jax's changelog.

JAX 0.6.2 (June 17, 2025)

  • New features:

    • Added {func}jax.tree.broadcast which implements a pytree prefix broadcasting helper.

  • Changes

    • The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12.

JAX 0.6.1 (May 21, 2025)

  • New features:

    • Added {func}jax.lax.axis_size which returns the size of the mapped axis given its name.

  • Changes

    • Additional checking for the versions of CUDA package dependencies was re-enabled, having been accidentally disabled in a previous release.
    • JAX nightly packages are now published to artifact registry. To install these packages, see the JAX installation guide.
    • jax.sharding.PartitionSpec no longer inherits from a tuple.
    • jax.ShapeDtypeStruct is immutable now. Please use .update method to update your ShapeDtypeStruct instead of doing in-place updates.

  • Deprecations

    • jax.custom_derivatives.custom_jvp_call_jaxpr_p is deprecated, and will be removed in JAX v0.7.0.

JAX 0.6.0 (April 16, 2025)

  • Breaking changes

    • {func}jax.numpy.array no longer accepts None. This behavior was deprecated since November 2023 and is now removed.
    • Removed the config.jax_data_dependent_tracing_fallback config option, which was added temporarily in v0.4.36 to allow users to opt out of the new "stackless" tracing machinery.
    • Removed the config.jax_eager_pmap config option.
    • Disallow the calling of lower and trace AOT APIs on the result of jax.jit if there have been subsequent wrappers applied. Previously this worked, but silently ignored the wrappers. The workaround is to apply jax.jit last among the wrappers, and similarly for jax.pmap. See {jax-issue}[#27873](https://github.com/jax-ml/jax/issues/27873).
    • The cuda12_pip extra for jax has been removed; use pip install jax[cuda12] instead.

  • Changes

    • The minimum CuDNN version is v9.8.
    • JAX is now built using CUDA 12.8. All versions of CUDA 12.1 or newer remain supported.

... (truncated)

Commits
  • 1ad05bb Add LLVM patch to fix AVX512 codegeneration problem
  • 8f81490 Prepare for JAX release 0.6.2
  • e4de90e Update XLA dependency to use revision
  • 02688e1 [Pallas][Mosaic GPU] Enable collective MMA from TMEM.
  • dc9ef61 Merge pull request #29410 from DanisNone:nn-type
  • 353e7fa Merge pull request #29516 from jakevdp:enable-x64-warning
  • 3d37b0d Merge pull request #29504 from vfdev-5:tsan-ft-removed-fixed-suppression
  • c2cc9f9 [pallas] AbstractMemoryRef now implements all functional update methods via...
  • 0fd0821 [Pallas TPU] Add flag to enable using registers to keep track of slot info
  • f22896a jax.experimental.enable_x64: add warning to docstring
  • Additional commits viewable in compare view

@dependabot dependabot Bot added dependencies Pull requests that update a dependency file Python labels Dec 15, 2025
@dependabot dependabot Bot force-pushed the dependabot/pip/devel/jax-0.6.2 branch from 8aada41 to 4a5a77f Compare December 15, 2025 22:49
@njzjz njzjz added this pull request to the merge queue Dec 16, 2025
@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Dec 16, 2025
@njzjz njzjz enabled auto-merge December 17, 2025 17:32
@codecov
Copy link
Copy Markdown

codecov Bot commented Dec 17, 2025 •
edited
Loading

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.45%. Comparing base (eab3419) to head (f03491d).

Additional details and impacted files 
@@            Coverage Diff             @@
##           master    #5102      +/-   ##
==========================================
- Coverage   82.47%   82.45%   -0.03%     
==========================================
  Files         825      825              
  Lines       87726    87727       +1     
  Branches     4206     4206              
==========================================
- Hits        72355    72335      -20     
- Misses      14088    14110      +22     
+ Partials     1283     1282       -1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@njzjz njzjz added this pull request to the merge queue Dec 17, 2025
@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Dec 17, 2025
@wanghan-iapcm wanghan-iapcm added this pull request to the merge queue Dec 23, 2025
@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Dec 23, 2025
njzjz
njzjz requested changes Dec 23, 2025
View reviewed changes
Copy link
Copy Markdown
Member

@njzjz njzjz left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need TensorFlow to bump to a specific version as well. This is the limitation of jax2tf.

njzjz
njzjz approved these changes Dec 24, 2025
View reviewed changes
Copy link
Copy Markdown
Member

@njzjz njzjz left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Try 2.19.1

@njzjz njzjz enabled auto-merge December 24, 2025 17:54
@njzjz njzjz added this pull request to the merge queue Dec 24, 2025
@github-merge-queue github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Dec 24, 2025
@dependabot @github
Copy link
Copy Markdown
Contributor Author

dependabot Bot commented on behalf of github Jan 6, 2026

A newer version of jax exists, but since this PR has been edited by someone other than Dependabot I haven't updated it. You'll get a PR for the updated version as normal once this PR is merged.

@njzjz njzjz marked this pull request as draft February 2, 2026 20:17
@njzjz
Copy link
Copy Markdown
Member

njzjz commented May 9, 2026

@dependabot recreate

@dependabot
build(deps-dev): bump jax from 0.5.0 to 0.6.2
f03491d 
Bumps [jax](https://github.com/jax-ml/jax) from 0.5.0 to 0.6.2.
- [Release notes](https://github.com/jax-ml/jax/releases)
- [Changelog](https://github.com/jax-ml/jax/blob/main/CHANGELOG.md)
- [Commits](jax-ml/jax@jax-v0.5.0...jax-v0.6.2)

---
updated-dependencies:
- dependency-name: jax
  dependency-version: 0.6.2
  dependency-type: direct:development
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
@dependabot dependabot Bot force-pushed the dependabot/pip/devel/jax-0.6.2 branch from c988731 to f03491d Compare May 9, 2026 16:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@njzjz njzjz njzjz approved these changes

Assignees

No one assigned

Labels

dependencies Pull requests that update a dependency file Python

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@njzjz